package mem

import (
	"errors"
	"fmt"
	"sync"
)

var (
	bufPoolInstance *BufPool
	once            sync.Once
)

const (
	// 总内存池大小限制 5GB
	EXTRA_MEM_LIMIT int = 5 * m1K * m1K
)

const (
	m1K int = (1 << 10) * (1 << iota) * (1 << iota)
	m4K
	m16K
	m64K
	m256K
	m1M
	m4M
	m8M = m4M << 1
)

type Pool map[int]*Buf

// BufPool 内存池
type BufPool struct {
	Pool     Pool   // 所有Buf的map集合句柄
	TotalMem uint64 // 总buffer池内存大小，单位KB
	PoolLock sync.RWMutex
}

// GetBufPoolInstance 获取单例模式的内存池实例
func GetBufPoolInstance() *BufPool {
	once.Do(func() {
		bufPoolInstance = new(BufPool)
		bufPoolInstance.Pool = make(map[int]*Buf)
		bufPoolInstance.TotalMem = 0
		bufPoolInstance.initPool()
	})
	return bufPoolInstance
}

// Alloc 申请一块内存空间，内存池存在则直接返回，不存在则新建
func (bp *BufPool) Alloc(N int) (*Buf, error) {
	// 1. 找到最匹配 N 的一个 hash 组
	var index int
	if N <= m4K {
		index = m4K
	} else if N <= m16K {
		index = m16K
	} else if N <= m64K {
		index = m64K
	} else if N <= m256K {
		index = m256K
	} else if N <= m1M {
		index = m1M
	} else if N <= m4M {
		index = m4M
	} else if N <= m8M {
		index = m8M
	} else {
		return nil, errors.New("Alloc size Too Large!")
	}

	bp.PoolLock.Lock()
	// 2. 该组没有，额外申请
	if bp.Pool[index] == nil {
		if bp.TotalMem+uint64(index/1024) >= uint64(EXTRA_MEM_LIMIT) {
			return nil, errors.New("Already use too many memory!")
		}

		newBuf := NewBuf(index)
		bp.TotalMem += uint64(index / 1024)
		bp.PoolLock.Unlock()
		return newBuf, nil
	}

	// 3. 该组有返回即可
	targetBuf := bp.Pool[index]
	bp.Pool[index] = targetBuf.Next
	bp.TotalMem -= uint64(index / 1024)
	bp.PoolLock.Unlock()
	targetBuf.Next = nil
	return targetBuf, nil
}

// Revert 归还内存块 Buf
func (bp *BufPool) Revert(buf *Buf) error {
	index := buf.Capacity
	buf.Clear()

	bp.PoolLock.Lock()
	// 1. Pool 中不存在容量为 index 的内存块
	if _, ok := bp.Pool[index]; !ok {
		errStr := fmt.Sprintf("Capacity %d not in BufPool\n", index)
		return errors.New(errStr)
	}

	// 2. 存在，头插法插回链表头部
	buf.Next = bp.Pool[index]
	bp.Pool[index] = buf
	bp.TotalMem += uint64(index / 1024)
	bp.PoolLock.Unlock()

	return nil
}

// initPool 初始化内存池
func (bp *BufPool) initPool() {
	// 4KB 开辟 5000 个，大约 20MB
	bp.makeBufList(m4K, 5000)
	// 16KB 开辟 1000 个，大约 16MB
	bp.makeBufList(m16K, 1000)
	// 64KB 开辟 500 个，大约 32MB
	bp.makeBufList(m64K, 500)
	// 256KB 开辟 200 个，大约 50MB
	bp.makeBufList(m256K, 200)
	// 1MB 开辟 50 个，大约 50MB
	bp.makeBufList(m1M, 50)
	// 4MB 开辟 20 个，大约 80MB
	bp.makeBufList(m4M, 20)
	// 8MB 开辟 10 个，大约 80MB
	bp.makeBufList(m8M, 10)
}

// makeBufList 创建 num 个 cap 容量大小的内存块
func (bp *BufPool) makeBufList(cap, num int) {
	bp.Pool[cap] = NewBuf(cap)

	var prev *Buf
	prev = bp.Pool[cap]
	for i := 1; i < num; i++ {
		prev.Next = NewBuf(cap)
		prev = prev.Next
	}
	bp.TotalMem += (uint64(cap) / 1024) * uint64(num)
}
